"""
检索器模块

实现不同类型的检索器和检索功能。
"""

from typing import List, Optional, Dict, Any
from abc import ABC, abstractmethod

from llama_index.core import VectorStoreIndex
from llama_index.core.schema import NodeWithScore, QueryBundle, QueryType
from llama_index.core.retrievers import BaseRetriever
# 尝试导入BaseNodePostprocessor，如果失败则使用try-except
try:
    from llama_index.core.postprocessor import BaseNodePostprocessor
except ImportError:
    # 在较新版本的LlamaIndex中，这个类可能位于不同的位置
    try:
        from llama_index.core.schema import BaseNodePostprocessor
    except ImportError:
        # 如果都找不到，定义一个基础类
        from abc import ABC
        class BaseNodePostprocessor(ABC):
            pass
from llama_index.core.query_engine import RetrieverQueryEngine

from .strategies import RetrievalStrategy, RetrievalConfig


class CustomRetriever(BaseRetriever):
    """
    自定义检索器

    支持多种检索策略和后处理管道。
    """

    def __init__(
        self,
        index: VectorStoreIndex,
        strategy: RetrievalStrategy,
        postprocessors: Optional[List[BaseNodePostprocessor]] = None
    ):
        """
        初始化自定义检索器

        Args:
            index: 向量索引
            strategy: 检索策略
            postprocessors: 后处理器列表
        """
        self._index = index
        self._strategy = strategy
        self._postprocessors = postprocessors or []
        self._retriever = self._create_base_retriever()
        super().__init__()

    def _create_base_retriever(self) -> BaseRetriever:
        """创建基础检索器"""
        retriever_kwargs = self._strategy.get_retriever_kwargs()

        # 根据策略配置创建不同类型的检索器
        if self._strategy.config.vector_store_query_mode == "hybrid":
            # 混合检索（向量+稀疏）
            return self._index.as_retriever(
                **retriever_kwargs,
                vector_store_query_mode="hybrid",
                alpha=0.5  # 向量和稀疏检索的权重
            )
        elif self._strategy.config.vector_store_query_mode == "mmr":
            # MMR（最大边际相关性）检索
            return self._index.as_retriever(
                **retriever_kwargs,
                vector_store_query_mode="mmr"
            )
        else:
            # 标准向量检索
            return self._index.as_retriever(**retriever_kwargs)

    def _retrieve(self, query_bundle: QueryBundle) -> List[NodeWithScore]:
        """
        执行检索

        Args:
            query_bundle: 查询包

        Returns:
            List[NodeWithScore]: 检索结果节点
        """
        # 验证查询
        if not self._strategy.validate_query(query_bundle.query_str):
            return []

        # 执行基础检索
        nodes = self._retriever.retrieve(query_bundle)

        # 应用相似度截断
        if self._strategy.config.similarity_cutoff:
            nodes = [
                node for node in nodes
                if node.score and node.score >= self._strategy.config.similarity_cutoff
            ]

        # 应用后处理器
        for postprocessor in self._postprocessors:
            nodes = postprocessor.postprocess_nodes(
                nodes, query_bundle=query_bundle
            )

        return nodes

    def add_postprocessor(self, postprocessor: BaseNodePostprocessor) -> None:
        """
        添加后处理器

        Args:
            postprocessor: 后处理器实例
        """
        self._postprocessors.append(postprocessor)

    def remove_postprocessor(self, postprocessor_class: type) -> None:
        """
        移除指定类型的后处理器

        Args:
            postprocessor_class: 后处理器类
        """
        self._postprocessors = [
            p for p in self._postprocessors
            if not isinstance(p, postprocessor_class)
        ]

    def get_retrieval_stats(self) -> Dict[str, Any]:
        """
        获取检索统计信息

        Returns:
            Dict[str, Any]: 统计信息
        """
        return {
            "strategy": type(self._strategy).__name__,
            "config": self._strategy.config.to_dict(),
            "postprocessor_count": len(self._postprocessors),
            "postprocessor_types": [type(p).__name__ for p in self._postprocessors]
        }


class MultiRetriever:
    """
    多检索器聚合

    支持同时使用多个检索器并聚合结果。
    """

    def __init__(
        self,
        retrievers: List[BaseRetriever],
        weights: Optional[List[float]] = None,
        aggregation_method: str = "weighted_average"  # "weighted_average", "union", "intersection"
    ):
        """
        初始化多检索器

        Args:
            retrievers: 检索器列表
            weights: 权重列表
            aggregation_method: 聚合方法
        """
        if not retrievers:
            raise ValueError("至少需要一个检索器")

        self.retrievers = retrievers
        self.weights = weights or [1.0] * len(retrievers)
        self.aggregation_method = aggregation_method

        if len(self.weights) != len(self.retrievers):
            raise ValueError("权重数量必须与检索器数量相同")

        # 归一化权重
        total_weight = sum(self.weights)
        self.weights = [w / total_weight for w in self.weights]

    def retrieve(self, query: QueryType) -> List[NodeWithScore]:
        """
        执行多检索器检索

        Args:
            query: 查询对象

        Returns:
            List[NodeWithScore]: 聚合后的检索结果
        """
        # 标准化查询
        if not isinstance(query, QueryBundle):
            query = QueryBundle(query_str=str(query))

        # 执行所有检索器
        all_results = []
        for i, retriever in enumerate(self.retrievers):
            try:
                results = retriever.retrieve(query)
                # 应用权重
                weight = self.weights[i]
                weighted_results = []
                for node in results:
                    weighted_node = NodeWithScore(
                        node=node.node,
                        score=node.score * weight if node.score else weight
                    )
                    weighted_results.append(weighted_node)
                all_results.append(weighted_results)
            except Exception as e:
                print(f"检索器 {i} 执行失败: {e}")
                continue

        # 聚合结果
        if self.aggregation_method == "union":
            return self._union_aggregate(all_results)
        elif self.aggregation_method == "intersection":
            return self._intersection_aggregate(all_results)
        else:  # weighted_average
            return self._weighted_average_aggregate(all_results)

    def _union_aggregate(self, all_results: List[List[NodeWithScore]]) -> List[NodeWithScore]:
        """联合聚合"""
        # 合并所有结果并去重
        seen_nodes = set()
        aggregated_results = []

        for results in all_results:
            for node in results:
                node_id = node.node.node_id if hasattr(node.node, 'node_id') else str(hash(node.node.text))
                if node_id not in seen_nodes:
                    seen_nodes.add(node_id)
                    aggregated_results.append(node)

        # 按分数排序
        aggregated_results.sort(key=lambda x: x.score or 0, reverse=True)
        return aggregated_results

    def _intersection_aggregate(self, all_results: List[List[NodeWithScore]]) -> List[NodeWithScore]:
        """交集聚合"""
        if not all_results:
            return []

        # 找出所有结果中都出现的节点
        node_scores = {}
        for i, results in enumerate(all_results):
            for node in results:
                node_id = node.node.node_id if hasattr(node.node, 'node_id') else str(hash(node.node.text))
                if node_id not in node_scores:
                    node_scores[node_id] = {
                        'node': node.node,
                        'scores': [],
                        'present_in': 0
                    }
                node_scores[node_id]['scores'].append(node.score or 0)
                node_scores[node_id]['present_in'] += 1

        # 只保留在所有检索器中都出现的节点
        intersection_results = []
        for node_id, node_info in node_scores.items():
            if node_info['present_in'] == len(all_results):
                # 计算平均分数
                avg_score = sum(node_info['scores']) / len(node_info['scores'])
                intersection_results.append(
                    NodeWithScore(node=node_info['node'], score=avg_score)
                )

        # 按分数排序
        intersection_results.sort(key=lambda x: x.score, reverse=True)
        return intersection_results

    def _weighted_average_aggregate(self, all_results: List[List[NodeWithScore]]) -> List[NodeWithScore]:
        """加权平均聚合"""
        node_scores = {}

        for results in all_results:
            for node in results:
                node_id = node.node.node_id if hasattr(node.node, 'node_id') else str(hash(node.node.text))
                if node_id not in node_scores:
                    node_scores[node_id] = {
                        'node': node.node,
                        'total_score': 0,
                        'count': 0
                    }
                node_scores[node_id]['total_score'] += node.score or 0
                node_scores[node_id]['count'] += 1

        # 计算加权平均分数
        aggregated_results = []
        for node_id, node_info in node_scores.items():
            avg_score = node_info['total_score'] / node_info['count']
            aggregated_results.append(
                NodeWithScore(node=node_info['node'], score=avg_score)
            )

        # 按分数排序
        aggregated_results.sort(key=lambda x: x.score, reverse=True)
        return aggregated_results

    def get_retriever_info(self) -> Dict[str, Any]:
        """
        获取检索器信息

        Returns:
            Dict[str, Any]: 检索器信息
        """
        return {
            "retriever_count": len(self.retrievers),
            "weights": self.weights,
            "aggregation_method": self.aggregation_method,
            "retriever_types": [type(r).__name__ for r in self.retrievers]
        }


class AdaptiveRetriever:
    """
    自适应检索器

    根据查询特征自动选择最适合的检索策略。
    """

    def __init__(self, index: VectorStoreIndex):
        """
        初始化自适应检索器

        Args:
            index: 向量索引
        """
        self.index = index
        self.retrievers = {}
        self._setup_retrievers()

    def _setup_retrievers(self) -> None:
        """设置不同策略的检索器"""
        from .strategies import (
            RetrievalMode, RetrievalStrategyFactory,
            StandardRetrievalStrategy, SimilarityRetrievalStrategy,
            HybridRetrievalStrategy, SemanticRetrievalStrategy
        )

        # 创建不同策略的检索器
        strategies = {
            RetrievalMode.STANDARD: StandardRetrievalStrategy(),
            RetrievalMode.SIMILARITY: SimilarityRetrievalStrategy(),
            RetrievalMode.HYBRID: HybridRetrievalStrategy(),
            RetrievalMode.SEMANTIC: SemanticRetrievalStrategy()
        }

        for mode, strategy in strategies.items():
            self.retrievers[mode] = CustomRetriever(self.index, strategy)

    def retrieve(self, query: QueryType) -> List[NodeWithScore]:
        """
        自适应检索

        Args:
            query: 查询对象

        Returns:
            List[NodeWithScore]: 检索结果
        """
        # 分析查询特征
        query_features = self._analyze_query(query)

        # 选择最适合的策略
        best_mode = self._select_best_strategy(query_features)

        # 使用选定的检索器
        retriever = self.retrievers[best_mode]
        return retriever.retrieve(query)

    def _analyze_query(self, query: QueryType) -> Dict[str, Any]:
        """分析查询特征"""
        query_str = str(query) if not isinstance(query, QueryBundle) else query.query_str

        features = {
            'length': len(query_str),
            'word_count': len(query_str.split()),
            'has_numbers': any(c.isdigit() for c in query_str),
            'is_question': any(q in query_str.lower() for q in ['什么', '怎么', '为什么', 'what', 'how', 'why']),
            'is_short': len(query_str) < 20,
            'is_long': len(query_str) > 100,
            'has_special_chars': any(not c.isalnum() and not c.isspace() for c in query_str)
        }

        return features

    def _select_best_strategy(self, features: Dict[str, Any]) -> str:
        """选择最佳检索策略"""
        # 简单的策略选择逻辑
        if features['is_short']:
            return 'hybrid'  # 短查询使用混合检索提高召回率
        elif features['has_numbers']:
            return 'similarity'  # 包含数字的查询使用精确相似度检索
        elif features['is_long']:
            return 'semantic'  # 长查询使用语义检索
        elif features['is_question']:
            return 'semantic'  # 问题查询使用语义检索
        else:
            return 'standard'  # 默认使用标准检索

    def add_custom_retriever(self, mode: str, retriever: BaseRetriever) -> None:
        """
        添加自定义检索器

        Args:
            mode: 检索模式名称
            retriever: 检索器实例
        """
        self.retrievers[mode] = retriever

    def remove_retriever(self, mode: str) -> None:
        """
        移除检索器

        Args:
            mode: 检索模式名称
        """
        if mode in self.retrievers:
            del self.retrievers[mode]

    def get_available_strategies(self) -> List[str]:
        """
        获取可用的检索策略

        Returns:
            List[str]: 策略列表
        """
        return list(self.retrievers.keys())